{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RAGatouille + Instructor: Finetuning ColBERT(v2) with no annotated data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the [previous example](https://github.com/bclavie/RAGatouille/blob/main/examples/02-basic_training.ipynb), we experimented with training our own ColBERT model. The process was straightforward, but our resulting model wasn't very useful: we didn't have any of the annotated data needed, so we created meaningless `[query, relevant_passage]` pairs that the model couldn't learn anything from!\n",
    "\n",
    "Getting annotated data is expensive! Thankfully, the literature in retrieval has recently shown that synthetic data can yield similar, if not better, performance when fine-tuning retrieval models. This means we can fine-tune to our target domain without **needing pricey and time-consuming annotations**.\n",
    "\n",
    "In this tutorial, we'll show how easily we can leverage Jason Wei's [instructor library](https://github.com/jxnl/instructor) for structured extraction with OpenAI's functional calling API to generate meaningful query-passage pairs.\n",
    "\n",
    "First, we'll need to install the required dependencies:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33mWARNING: Ignoring invalid distribution -orch (/opt/conda/lib/python3.10/site-packages)\u001b[0m\u001b[33m\n",
      "\u001b[0mRequirement already satisfied: openai in /opt/conda/lib/python3.10/site-packages (1.6.1)\n",
      "Requirement already satisfied: instructor in /opt/conda/lib/python3.10/site-packages (0.4.5)\n",
      "Requirement already satisfied: anyio<5,>=3.5.0 in /opt/conda/lib/python3.10/site-packages (from openai) (3.7.1)\n",
      "Requirement already satisfied: distro<2,>=1.7.0 in /opt/conda/lib/python3.10/site-packages (from openai) (1.8.0)\n",
      "Requirement already satisfied: httpx<1,>=0.23.0 in /opt/conda/lib/python3.10/site-packages (from openai) (0.26.0)\n",
      "Requirement already satisfied: pydantic<3,>=1.9.0 in /opt/conda/lib/python3.10/site-packages (from openai) (2.5.3)\n",
      "Requirement already satisfied: sniffio in /opt/conda/lib/python3.10/site-packages (from openai) (1.3.0)\n",
      "Requirement already satisfied: tqdm>4 in /opt/conda/lib/python3.10/site-packages (from openai) (4.66.1)\n",
      "Requirement already satisfied: typing-extensions<5,>=4.7 in /opt/conda/lib/python3.10/site-packages (from openai) (4.9.0)\n",
      "Requirement already satisfied: aiohttp<4.0.0,>=3.9.1 in /opt/conda/lib/python3.10/site-packages (from instructor) (3.9.1)\n",
      "Requirement already satisfied: docstring-parser<0.16,>=0.15 in /opt/conda/lib/python3.10/site-packages (from instructor) (0.15)\n",
      "Requirement already satisfied: rich<14.0.0,>=13.7.0 in /opt/conda/lib/python3.10/site-packages (from instructor) (13.7.0)\n",
      "Requirement already satisfied: typer<0.10.0,>=0.9.0 in /opt/conda/lib/python3.10/site-packages (from instructor) (0.9.0)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /opt/conda/lib/python3.10/site-packages (from aiohttp<4.0.0,>=3.9.1->instructor) (23.1.0)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /opt/conda/lib/python3.10/site-packages (from aiohttp<4.0.0,>=3.9.1->instructor) (6.0.4)\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in /opt/conda/lib/python3.10/site-packages (from aiohttp<4.0.0,>=3.9.1->instructor) (1.9.4)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /opt/conda/lib/python3.10/site-packages (from aiohttp<4.0.0,>=3.9.1->instructor) (1.4.0)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in /opt/conda/lib/python3.10/site-packages (from aiohttp<4.0.0,>=3.9.1->instructor) (1.3.1)\n",
      "Requirement already satisfied: async-timeout<5.0,>=4.0 in /opt/conda/lib/python3.10/site-packages (from aiohttp<4.0.0,>=3.9.1->instructor) (4.0.3)\n",
      "Requirement already satisfied: idna>=2.8 in /opt/conda/lib/python3.10/site-packages (from anyio<5,>=3.5.0->openai) (3.6)\n",
      "Requirement already satisfied: exceptiongroup in /opt/conda/lib/python3.10/site-packages (from anyio<5,>=3.5.0->openai) (1.2.0)\n",
      "Requirement already satisfied: certifi in /opt/conda/lib/python3.10/site-packages (from httpx<1,>=0.23.0->openai) (2023.11.17)\n",
      "Requirement already satisfied: httpcore==1.* in /opt/conda/lib/python3.10/site-packages (from httpx<1,>=0.23.0->openai) (1.0.2)\n",
      "Requirement already satisfied: h11<0.15,>=0.13 in /opt/conda/lib/python3.10/site-packages (from httpcore==1.*->httpx<1,>=0.23.0->openai) (0.14.0)\n",
      "Requirement already satisfied: annotated-types>=0.4.0 in /opt/conda/lib/python3.10/site-packages (from pydantic<3,>=1.9.0->openai) (0.6.0)\n",
      "Requirement already satisfied: pydantic-core==2.14.6 in /opt/conda/lib/python3.10/site-packages (from pydantic<3,>=1.9.0->openai) (2.14.6)\n",
      "Requirement already satisfied: markdown-it-py>=2.2.0 in /opt/conda/lib/python3.10/site-packages (from rich<14.0.0,>=13.7.0->instructor) (3.0.0)\n",
      "Requirement already satisfied: pygments<3.0.0,>=2.13.0 in /opt/conda/lib/python3.10/site-packages (from rich<14.0.0,>=13.7.0->instructor) (2.17.2)\n",
      "Requirement already satisfied: click<9.0.0,>=7.1.1 in /opt/conda/lib/python3.10/site-packages (from typer<0.10.0,>=0.9.0->instructor) (8.1.7)\n",
      "Requirement already satisfied: mdurl~=0.1 in /opt/conda/lib/python3.10/site-packages (from markdown-it-py>=2.2.0->rich<14.0.0,>=13.7.0->instructor) (0.1.2)\n",
      "\u001b[33mWARNING: Ignoring invalid distribution -orch (/opt/conda/lib/python3.10/site-packages)\u001b[0m\u001b[33m\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "!pip install openai instructor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To continue with this tutorial, you'll need to set up your OpenAI api key as an environment variable. If you're not planning on sharing this notebook with anyone, you can do so by uncommenting the line below and filling it with your own API key before running the next cell!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "# os.environ['OPENAI_API_KEY'] = \"YOUR_KEY_HERE\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that our API key's set-up, we will load up our OpenAI client and set-up `instructor`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import instructor\n",
    "# If you're using llamaindex 0.10 or above, these need to be imported from llama_index.core instead!\n",
    "from llama_index import Document\n",
    "from llama_index.text_splitter import SentenceSplitter\n",
    "from openai import OpenAI\n",
    "from typing import List\n",
    "from pydantic import BaseModel, Field\n",
    "\n",
    "client = instructor.patch(OpenAI(api_key=os.environ[\"OPENAI_API_KEY\"]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Perfect, now all we need is some data! Just like in the last tutorial, we'll be getting a few wikipedia pages. First, let's define the helper function again:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import requests\n",
    "\n",
    "def get_wikipedia_page(title: str):\n",
    "    \"\"\"\n",
    "    Retrieve the full text content of a Wikipedia page.\n",
    "    \n",
    "    :param title: str - Title of the Wikipedia page.\n",
    "    :return: str - Full text content of the page as raw string.\n",
    "    \"\"\"\n",
    "    # Wikipedia API endpoint\n",
    "    URL = \"https://en.wikipedia.org/w/api.php\"\n",
    "\n",
    "    # Parameters for the API request\n",
    "    params = {\n",
    "        \"action\": \"query\",\n",
    "        \"format\": \"json\",\n",
    "        \"titles\": title,\n",
    "        \"prop\": \"extracts\",\n",
    "        \"explaintext\": True,\n",
    "    }\n",
    "\n",
    "    # Custom User-Agent header to comply with Wikipedia's best practices\n",
    "    headers = {\n",
    "        \"User-Agent\": \"RAGatouille_tutorial/0.0.1 (ben@clavie.eu)\"\n",
    "    }\n",
    "\n",
    "    response = requests.get(URL, params=params, headers=headers)\n",
    "    data = response.json()\n",
    "\n",
    "    # Extracting page content\n",
    "    page = next(iter(data['query']['pages'].values()))\n",
    "    return page['extract'] if 'extract' in page else None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's load up a few pages and pre-process them into smaller chunks using `CorpusProcessor`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "my_full_corpus = [get_wikipedia_page(\"Hayao_Miyazaki\")]\n",
    "my_full_corpus += [get_wikipedia_page(\"Studio_Ghibli\")]\n",
    "my_full_corpus += [get_wikipedia_page(\"Toei_Animation\")]\n",
    "\n",
    "from ragatouille.data import CorpusProcessor, llama_index_sentence_splitter\n",
    "\n",
    "corpus_processor = CorpusProcessor(document_splitter_fn=llama_index_sentence_splitter)\n",
    "documents = corpus_processor.process_corpus(my_full_corpus, chunk_size=180)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "How we proceed next would depend on your primary use-case. In a lot of situations, you may have a lot of documents, but think only some of them would be useful for user queries.\n",
    "\n",
    "We're going to make the same assumption here, so we'll randomly choose a few chunks as \"relevant passages\", for which we'll generate queries, and the others will serve as our full corpus to generate negatives from:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import random\n",
    "\n",
    "random.seed(42)\n",
    "relevant_documents = random.sample(documents, 32)\n",
    "irrelevant_documents = list(set(documents) - set(relevant_documents))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating our synthetic query set\n",
    "\n",
    "Now that we have relevant our list of documents, we need to generate queries for them to be able to train our model! Unlike in the last tutorial, we want actual, useful queries that will produce a model we can use.\n",
    "\n",
    "To do so, we'll define a `pydantic` model for `instructor` to use when calling OpenAI:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class QueryForPassage(BaseModel):\n",
    "    hypothetical_questions: List[str] = Field(\n",
    "        default_factory=list,\n",
    "        description=\"A wide variety of hypothetical questions that this document could answer.\",\n",
    "    )\n",
    "    hypothetical_queries: List[str] = Field(\n",
    "        default_factory=list,\n",
    "        description=\"A wide variety of hypothetical queries that this document would be relevant to, in the context of a search engine or a retrieval pipeline.\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We're defining two styles of queries here: the variety might be helpful in preparing the model for different keywords the users may use in their queries! Feel free to experiment with the descriptions, add your own, or even request that the model also assigns a few keywords to the document.\n",
    "\n",
    "For this example, we'll only run it on the one passage so it doesn't get too wordy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Document: \n",
      "In April 1984, Miyazaki opened his own office in Suginami Ward, naming it Nibariki.\n",
      "\n",
      "\n",
      "=== Studio Ghibli ===\n",
      "\n",
      "\n",
      "==== Early films (1985–1996) ====\n",
      "In June 1985, Miyazaki, Takahata, Tokuma and Suzuki founded the animation production company Studio Ghibli, with funding from Tokuma Shoten. Studio Ghibli's first film, Laputa: Castle in the Sky (1986), employed the same production crew of Nausicaä. Miyazaki's designs for the film's setting were inspired by Greek architecture and \"European urbanistic templates\".\n",
      "Generated queries: \n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'hypothetical_questions': [\"What was Studio Ghibli's first film?\",\n",
       "  \"Where was Miyazaki's office located when he opened it?\",\n",
       "  'Who were the founders of Studio Ghibli?',\n",
       "  'What year was Studio Ghibli founded?',\n",
       "  \"What inspired Miyazaki's designs for the setting of Laputa: Castle in the Sky?\",\n",
       "  'Which production crew did Laputa: Castle in the Sky employ?'],\n",
       " 'hypothetical_queries': ['First film produced by Studio Ghibli',\n",
       "  \"Location of Miyazaki's office in 1984\",\n",
       "  'Founders of Studio Ghibli',\n",
       "  'Year Studio Ghibli was established',\n",
       "  \"Inspirations for Miyazaki's film Laputa\",\n",
       "  'Production crew for Laputa: Castle in the Sky']}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "candidate_queries = []\n",
    "for doc in relevant_documents:\n",
    "    candidate = client.chat.completions.create(\n",
    "        model=\"gpt-4-1106-preview\",\n",
    "        response_model=QueryForPassage,\n",
    "        messages=[\n",
    "            {\n",
    "                \"role\": \"system\",\n",
    "                \"content\": \"\"\"You are an expert AI assisting us in creating a high quality, diverse synthetic dataset to train Information Retrieval models. \n",
    "Your role is to analyse the document chunk given to you and provide us with high quality potential queries.\"\"\",\n",
    "            },\n",
    "            {\"role\": \"user\", \"content\": doc},\n",
    "        ],\n",
    "    )\n",
    "    candidate_queries.append(candidate)\n",
    "    break\n",
    "\n",
    "print(\"Document: \")\n",
    "print(relevant_documents[0])\n",
    "print(\"Generated queries: \")\n",
    "candidate_queries[0].model_dump()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Obviously, in a real setting, you'd generate queries for many, many more of your documents to populate your training data! Feel free to experiment with the description you use to the model, and even sway it towards certain topics if you think those are particualrly important to your users.\n",
    "\n",
    "Now that we've generated the queries, we need to create training pairs. In the example below, we arbitrarily decide that we want to keep two potential query for each relevant document: one from `hypothetical_queries` and one from `hypothetical_questions`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "pairs = []\n",
    "num_questions = 1\n",
    "num_queries = 1\n",
    "random.seed(42)\n",
    "\n",
    "for candidates, doc in zip(candidate_queries, relevant_documents):\n",
    "    candidates = candidates.model_dump()\n",
    "    queries = random.sample(candidates['hypothetical_questions'], num_questions)\n",
    "    queries += random.sample(candidates['hypothetical_queries'], num_queries)\n",
    "    for q in queries:\n",
    "        pairs.append([q, doc])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And just like that, we've generated annotated training examples in just a few lines of code, thanks to the magic of `pydantic` harnessed by `instructor` and OpenAI's function calling API!\n",
    "\n",
    "Now, we just need to get our trainer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from ragatouille import RAGTrainer\n",
    "\n",
    "trainer = RAGTrainer(model_name=\"GhibliColBERTv2.0\", pretrained_model_name=\"colbert-ir/colbertv2.0\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we'll get our trainer to prepare our training data. We'll pass it both the pairs we've created, as well as the full document corpus, so it can generate hard negatives via `SimpleMiner`:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading Hard Negative SimpleMiner dense embedding model BAAI/bge-small-en-v1.5...\n",
      "Building hard negative index for 156 documents...\n",
      "All documents embedded, now adding to index...\n",
      "save_index set to False, skipping saving hard negative index\n",
      "Hard negative index generated\n",
      "mining\n",
      "mining\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'./data/'"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.prepare_training_data(\n",
    "        raw_data = pairs,\n",
    "        all_documents = documents,\n",
    "        num_new_negatives = 10,\n",
    "        mine_hard_negatives= True,\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our trainer is now ready to `train()`, just like in the [previous tutorial](https://github.com/bclavie/RAGatouille/blob/main/examples/02-basic_training.ipynb). We won't actually train a model here, since we created just one training pair, but you've now got all you need to fine-tune your own ColBERT on your data, without any annotations!\n",
    "\n",
    "\n",
    "**You're now done with the RAGatouille 0.0.1 examples!**\n",
    "\n",
    "\n",
    "In a future tutorial, we'll demonstrate how to leverage [DSPy](https://github.com/stanfordnlp/dspy) to perform this kind of data-generation using open-source language models, in case you want to avoid relying on API calls!\n",
    "\n",
    "In future RAGatouille version, we'll further improve on this approach by supporting [UDAPDR](https://arxiv.org/abs/2303.00807), the state-of-the-art method to adapt retrievers to any domain, using LLMs."
   ]
  }
 ],
 "metadata": {
  "environment": {
   "kernel": "python3",
   "name": ".m114",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/:m114"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
